// src/random_text.cpp
#include "random_text.h"
#include <fstream>
#include <sstream>
#include <iostream>
#include <random>
#include <unordered_set>

// 构造函数
MarkovModel::MarkovModel(int state_size) : state_size_(state_size) {}

// 训练模型
bool MarkovModel::train(const std::string& file_path) {
    std::ifstream file(file_path, std::ios::in);
    if (!file.is_open()) {
        std::cerr << "无法打开文件: " << file_path << std::endl;
        return false;
    }

    std::string line;
    std::vector<std::string> tokens;
    std::vector<std::pair<std::string, std::string>> starting_keys;

    // 逐行读取文件，每行代表一个句子
    while (std::getline(file, line)) {
        if (line.empty()) continue;
        std::stringstream ss(line);
        std::vector<std::string> words;
        std::string word;
        while (ss >> word) {
            words.push_back(word);
        }

        if (words.size() < state_size_) continue; // 句子中词语不足

        // 将句子的前两个词作为起始键
        starting_keys.emplace_back(std::make_pair(words[0], words[1]));

        // 将所有词语添加到tokens向量中
        for (const auto& w : words) {
            tokens.push_back(w);
        }
    }

    if (tokens.size() < static_cast<size_t>(state_size_) + 1) {
        std::cerr << "语料库太小，无法训练模型。" << std::endl;
        return false;
    }

    // 构建转移表
    for (size_t i = 0; i + state_size_ < tokens.size(); ++i) {
        std::pair<std::string, std::string> key = std::make_pair(tokens[i], tokens[i + 1]);
        std::string next_word = tokens[i + 2];
        transitions_[key].push_back(next_word);
    }

    // 保存起始键
    this->starting_keys_ = starting_keys;

    file.close();
    return true;
}

// // 生成句子
// std::string MarkovModel::generate_sentence(int max_words, unsigned int seed) {
//     if (transitions_.empty() || starting_keys_.empty()) {
//         return "没有足够的训练数据来生成句子。";
//     }
//
//     // 初始化随机数生成器
//     std::mt19937 gen(seed ? seed : std::random_device{}());
//
//     // 尝试最多300次生成句子
//     std::string sentence;
//     for (int attempt = 0; attempt < 300; ++attempt) {
//         // 随机选择一个起始键
//         std::uniform_int_distribution<> dist_start(0, starting_keys_.size() - 1);
//         std::pair<std::string, std::string> current_key = starting_keys_[dist_start(gen)];
//
//         // 初始化句子
//         std::stringstream result;
//         result << current_key.first << current_key.second;
//
//         int words_generated = 2;
//
//         // 逐步生成后续词语
//         while (words_generated < max_words) {
//             auto it = transitions_.find(current_key);
//             if (it == transitions_.end() || it->second.empty()) {
//                 break;
//             }
//
//             // 随机选择下一个词
//             const std::vector<std::string>& possible_next = it->second;
//             std::uniform_int_distribution<> dist_word(0, possible_next.size() - 1);
//             std::string next_word = possible_next[dist_word(gen)];
//
//             // 添加到句子中
//             result << next_word;
//             words_generated++;
//
//             // 更新当前键
//             current_key = std::make_pair(current_key.second, next_word);
//         }
//
//         // 获取生成的句子
//         sentence = result.str();
//
//         // 检查句子是否以标点符号结尾
//         if (!sentence.empty()) {
//             std::string last_char_str = sentence.substr(sentence.size() - 1);
//             if (last_char_str == "。" || last_char_str == "！" || last_char_str == "？") {
//                 break;
//             }
//         }
//     }
//
//     if (sentence.empty()) {
//         return "生成失败，请重试。";
//     }
//     else {
//         return sentence;
//     }
// }


std::string MarkovModel::generate_sentence(int max_words, unsigned int seed) {
    if (transitions_.empty() || starting_keys_.empty()) {
        return "没有足够的训练数据来生成句子。";
    }

    // 初始化随机数生成器及相关分布对象，移到外层避免重复创建
    std::mt19937 gen(seed? seed : std::random_device{}());
    std::uniform_int_distribution<> dist_start(0, starting_keys_.size() - 1);
    std::unordered_set<std::string> punctuation_set = { "。", "！", "？" };  // 构建标点符号集合，方便判断

    // 尝试生成句子
    for (int attempt = 0; attempt < 300; ++attempt) {
        // 随机选择一个起始键
        std::pair<std::string, std::string> current_key = starting_keys_[dist_start(gen)];
        std::stringstream result;
        result << current_key.first << current_key.second;

        int words_generated = 2;
        bool sentence_ended = false;  // 标记句子是否已经以标点符号结尾

        // 逐步生成后续词语
        while (words_generated < max_words &&!sentence_ended) {
            auto it = transitions_.find(current_key);
            if (it == transitions_.end() || it->second.empty()) {
                break;
            }

            // 随机选择下一个词
            const std::vector<std::string>& possible_next = it->second;
            std::uniform_int_distribution<> dist_word(0, possible_next.size() - 1);
            std::string next_word = possible_next[dist_word(gen)];

            // 添加到句子中并立即检查是否以标点符号结尾
            result << next_word;
            words_generated++;
            if (punctuation_set.count(next_word) > 0) {  // 通过集合查找判断标点符号，更高效
                sentence_ended = true;
            }

            // 更新当前键
            current_key = std::make_pair(current_key.second, next_word);
        }

        std::string sentence = result.str();
        if (!sentence.empty() && sentence_ended) {
            return sentence;
        }
    }

    return "生成失败，请重试。";
}

// 生成随机句子的函数实现
std::string generate_random_sentence(MarkovModel& model, int max_words, unsigned int seed) {
    return model.generate_sentence(max_words, seed);
}
